﻿using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using EvoBoost.Core.Evolutionary;

namespace EvoBoost.Core
{
	internal class TreeChromosome<T> : IChromosome<TreeChromosome<T>>
	{
		private static readonly Random random = new Random();

		private readonly SplitBase<T> split;

		private readonly TreeNode<T> treeRoot;

		private bool isTreeDirty;

		private readonly ILossFunction lossFunction;

		private readonly int maxTreeRegions;

		private readonly int sampleCount;

		private readonly int minSamplesInNode;

		private readonly double adjustSplitterMutationProbability;

		public TreeChromosome(TreeChromosome<T> other)
		{
			Debug.Assert(other != null);

			this.split = other.split;
			this.sampleCount = other.sampleCount;
			this.lossFunction = other.lossFunction;
			this.maxTreeRegions = other.maxTreeRegions;
			this.minSamplesInNode = other.minSamplesInNode;
			this.adjustSplitterMutationProbability = other.adjustSplitterMutationProbability;
			
			this.treeRoot = other.treeRoot.CloneSubtree();
			this.isTreeDirty = other.isTreeDirty;
		}

		public TreeChromosome(
			SplitBase<T> split,
			ICollection<double> previousOutcomes,
			ICollection<Sample<T>> samples,
			ILossFunction lossFunction,
			int maxTreeRegions,
			int minSamplesInNode,
			double adjustSplitterMutationProbability)
			: this(split, previousOutcomes, samples, lossFunction, maxTreeRegions, minSamplesInNode, adjustSplitterMutationProbability, null)
		{
		}

		public TreeChromosome(
			SplitBase<T> split,
			ICollection<double> previousOutcomes,
			ICollection<Sample<T>> samples,
			ILossFunction lossFunction,
			int maxTreeRegions,
			int minSamplesInNode,
			double adjustSplitterMutationProbability,
			TreeNode<T> treeRoot)
		{
			Debug.Assert(split != null);
			Debug.Assert(minSamplesInNode >= 1);
			Debug.Assert(samples != null && samples.Count >= minSamplesInNode);
			Debug.Assert(previousOutcomes != null && previousOutcomes.Count == samples.Count);
			Debug.Assert(lossFunction != null);
			Debug.Assert(maxTreeRegions >= 1);
			Debug.Assert(adjustSplitterMutationProbability >= 0 && adjustSplitterMutationProbability <= 1);

			this.split = split;
			this.sampleCount = samples.Count;
			this.lossFunction = lossFunction;
			this.maxTreeRegions = maxTreeRegions;
			this.minSamplesInNode = minSamplesInNode;
			this.adjustSplitterMutationProbability = adjustSplitterMutationProbability;

			this.treeRoot = treeRoot ?? TreeNode<T>.CreateRoot(previousOutcomes, samples);
			this.isTreeDirty = true;
		}

		public double Fitness
		{
			get
			{
				this.RefreshTreeParametersIfNecessary();
				return -this.treeRoot.TotalLoss;
			}
		}

		public Tree<T> BuildTree()
		{
			this.RefreshTreeParametersIfNecessary();
			return new Tree<T>(this.treeRoot, this.split);
		}

		public void CrossoverWith(TreeChromosome<T> other)
		{
			// Extract all the nodes from trees
			List<TreeNode<T>> allNodes1 = new List<TreeNode<T>>();
			List<TreeNode<T>> allNodes2 = new List<TreeNode<T>>();
			this.treeRoot.ExtractSubtree(allNodes1, null, null);
			other.treeRoot.ExtractSubtree(allNodes2, null, null);

			// Build list of subtrees that can be swapped
			List<CrossoverPair> possibleCrossoverPairs = new List<CrossoverPair>();
			foreach (TreeNode<T> node1 in allNodes1)
			{
				// We don't want to swap roots
				if (node1 == this.treeRoot)
					continue;

				foreach (TreeNode<T> node2 in allNodes2)
				{
					// We don't want to swap roots or terminals
					if (node2 == other.treeRoot || (node1.Type == TreeNodeType.Terminal && node2.Type == TreeNodeType.Terminal))
						continue;

					// Check if subtree swap will not make new trees exceed max allowed terminal node count
					int terminalNodeDiff = node1.TerminalsInSubtree - node2.TerminalsInSubtree;
					if (this.treeRoot.TerminalsInSubtree - terminalNodeDiff <= this.maxTreeRegions &&
						other.treeRoot.TerminalsInSubtree + terminalNodeDiff <= other.maxTreeRegions)
					{
						possibleCrossoverPairs.Add(new CrossoverPair(node1, node2));
					}
				}
			}

			// No crossover available?
			if (possibleCrossoverPairs.Count == 0)
				return;

			// Select random crossover and perform it
			CrossoverPair selectedCrossover = possibleCrossoverPairs[random.Next(possibleCrossoverPairs.Count)];
			TreeNode<T>.SwapSubtrees(selectedCrossover.TreeNode1, selectedCrossover.TreeNode2);

			// Setup dirty flags
			this.isTreeDirty = true;
			other.isTreeDirty = true;
		}

		public TreeChromosome<T> Clone()
		{
			return new TreeChromosome<T>(this);
		}

		public TreeChromosome<T> CreateRandom()
		{
			TreeChromosome<T> clone = this.Clone();
			int mutationCount = random.Next(this.maxTreeRegions) + 1;	// TODO: change to something more reasonable
			for (int i = 0; i < mutationCount; ++i)
				clone.ApplyRandomMutation();

			return clone;
		}

		public void Mutate()
		{
			this.ApplyRandomMutation();
		}

		private void RefreshTreeParametersIfNecessary()
		{
			if (this.isTreeDirty)
			{
				this.treeRoot.RefreshParametersInSubtree(this.lossFunction, this.split, this.minSamplesInNode);
				this.isTreeDirty = false;
			}
		}

		private void ApplyRandomMutation()
		{
			// Extract terminal and splitter nodes from tree
			List<TreeNode<T>> terminalNodes = new List<TreeNode<T>>();
			List<TreeNode<T>> splitterNodes = new List<TreeNode<T>>();
			this.treeRoot.ExtractSubtree(null, terminalNodes, splitterNodes);

			// Find all possible mutations and save by type
			List<MutationDescription> splitTerminalMutations = new List<MutationDescription>();
			List<MutationDescription> adjustSplitterMutations = new List<MutationDescription>();
			List<MutationDescription> collapseSplitterMutations = new List<MutationDescription>();
			// Find all possible "split terminal" mutations
			if (this.treeRoot.TerminalsInSubtree < this.maxTreeRegions)	// Still can add a terminal node?
			{
				foreach (TreeNode<T> node in terminalNodes)
					splitTerminalMutations.Add(new MutationDescription(node, MutationType.SplitTerminal));
			}
			// Find all possible "adjust splitter" and "collapse splitter" mutations
			foreach (TreeNode<T> node in splitterNodes)
			{
				adjustSplitterMutations.Add(new MutationDescription(node, MutationType.AdjustSplitter));
				collapseSplitterMutations.Add(new MutationDescription(node, MutationType.CollapseSplitter));
			}

			// How many mutation groups are not empty
			int mutationGroups = 0;
			mutationGroups += splitTerminalMutations.Count > 0 ? 1 : 0;
			mutationGroups += adjustSplitterMutations.Count > 0 ? 1 : 0;
			mutationGroups += collapseSplitterMutations.Count > 0 ? 1 : 0;

			// No mutations available
			if (mutationGroups == 0)
				return;

			// Build list of possible mutations with probabilities
			int totalMutationCount = splitTerminalMutations.Count + adjustSplitterMutations.Count + collapseSplitterMutations.Count;
			List<MutationDescription> possibleMutations = new List<MutationDescription>(totalMutationCount);
			List<double> mutationProb = new List<double>(totalMutationCount);
			// Adjust splitter mutations
			if (adjustSplitterMutations.Count > 0)
			{
				double prob = totalMutationCount == adjustSplitterMutations.Count
								? 1.0 / adjustSplitterMutations.Count
								: this.adjustSplitterMutationProbability / adjustSplitterMutations.Count;
				for (int i = 0; i < adjustSplitterMutations.Count(); ++i)
					mutationProb.Add(prob);
				possibleMutations.AddRange(adjustSplitterMutations);
			}
			// Split terminal mutations
			if (splitTerminalMutations.Count > 0)
			{
				double localSpread = collapseSplitterMutations.Count > 0 ? 0.5 : 1;
				double globalSpread = adjustSplitterMutations.Count > 0 ? 1.0 - this.adjustSplitterMutationProbability : 1;
				double prob = localSpread * globalSpread / splitTerminalMutations.Count;
				for (int i = 0; i < splitTerminalMutations.Count(); ++i)
					mutationProb.Add(prob);
				possibleMutations.AddRange(splitTerminalMutations);
			}
			// Collapse splitter mutations
			if (collapseSplitterMutations.Count > 0)
			{
				double localSpread = splitTerminalMutations.Count > 0 ? 0.5 : 1;
				double globalSpread = adjustSplitterMutations.Count > 0 ? 1.0 - this.adjustSplitterMutationProbability : 1;
				double prob = localSpread * globalSpread / collapseSplitterMutations.Count;
				for (int i = 0; i < collapseSplitterMutations.Count(); ++i)
					mutationProb.Add(prob);
				possibleMutations.AddRange(collapseSplitterMutations);
			}

			// Choose and perform random mutation
			int mutationIndex = MathUtils.SampleFromDiscreteDistribution(mutationProb);
			MutationDescription randomMutation = possibleMutations[mutationIndex];
			switch (randomMutation.MutationType)
			{
				case MutationType.SplitTerminal:
					randomMutation.TreeNode.Split(this.split.GenerateRandomParameters());
					break;
				case MutationType.AdjustSplitter:
					randomMutation.TreeNode.Adjust(this.split.PerturbParameters(randomMutation.TreeNode.SplitParameters));
					break;
				case MutationType.CollapseSplitter:
					randomMutation.TreeNode.Collapse();
					break;
			}

			// Setup dirty flag
			this.isTreeDirty = true;
		}

		private class CrossoverPair
		{
			public TreeNode<T> TreeNode1 { get; private set; }

			public TreeNode<T> TreeNode2 { get; private set; }

			public CrossoverPair(TreeNode<T> treeNode1, TreeNode<T> treeNode2)
			{
				Debug.Assert(treeNode1 != null);
				Debug.Assert(treeNode2 != null);

				this.TreeNode1 = treeNode1;
				this.TreeNode2 = treeNode2;
			}
		}

		private class MutationDescription
		{
			public TreeNode<T> TreeNode { get; private set; }

			public MutationType MutationType { get; private set; }

			public MutationDescription(TreeNode<T> treeNode, MutationType mutationType)
			{
				Debug.Assert(treeNode != null);

				this.TreeNode = treeNode;
				this.MutationType = mutationType;
			}
		}
	}
}